Skip to content

Add BallTree from petal-neighbors for nearest neighbor search (Resolves #1)#34

Open
octoaide[bot] wants to merge 1 commit into
mainfrom
octoaide/implement-balltree
Open

Add BallTree from petal-neighbors for nearest neighbor search (Resolves #1)#34
octoaide[bot] wants to merge 1 commit into
mainfrom
octoaide/implement-balltree

Conversation

@octoaide
Copy link
Copy Markdown
Contributor

@octoaide octoaide Bot commented Dec 2, 2025

This PR implements the BallTree data structure from petal-neighbors as requested in issue #1, providing efficient nearest neighbor search capabilities to the pypetal project.

Closes #1

Summary

I've implemented the BallTree data structure from petal-neighbors 0.18.0 for the pypetal project. Here are the changes:

Files Modified

  1. Cargo.toml - Added petal-neighbors = "0.18" dependency and renamed the library from decomposition to pypetal to serve as the main module

  2. src/lib.rs - Updated to expose both decomposition and neighbors submodules through the main pypetal Python module

  3. src/decomposition.rs - Made the decomposition function public for use from lib.rs

  4. setup.py - Updated to use pypetal.pypetal as the Rust extension name

Files Created

  1. src/neighbors.rs - New module containing the BallTree wrapper with:
    • Constructor: BallTree(points) - Creates a ball tree from a 2D numpy array using Euclidean distance
    • Properties: n_samples, n_nodes
    • Methods:
      • query_nearest(point) - Returns the nearest neighbor index and distance
      • query(point, k) - Returns k nearest neighbor indices and distances
      • query_radius(point, radius) - Returns all neighbor indices within the given radius
    • Unit tests: 4 comprehensive tests for construction and all query methods

Python Usage Example

import numpy as np
from pypetal.pypetal import neighbors

# Create sample points
points = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]])

# Build the ball tree
tree = neighbors.BallTree(points)

# Query nearest neighbor
idx, dist = tree.query_nearest(np.array([0.1, 0.1]))

# Query k nearest neighbors
indices, distances = tree.query(np.array([0.0, 0.0]), k=2)

# Query all neighbors within radius
indices = tree.query_radius(np.array([0.0, 0.0]), 1.5)

All tests pass and clippy reports no warnings.

@octoaide octoaide Bot requested a review from msk December 2, 2025 23:37
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Add BallTree from petal-neighbors

0 participants